import os

from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_experimental.synthetic_data import create_data_generation_chain
from langchain_experimental.tabular_synthetic_data.openai import create_openai_data_generator
from langchain_experimental.tabular_synthetic_data.prompts import SYNTHETIC_FEW_SHOT_SUFFIX, SYNTHETIC_FEW_SHOT_PREFIX
from langchain_openai import ChatOpenAI
from pydantic import BaseModel

model = ChatOpenAI(model='gpt-4-turbo', temperature=0.8)

# experimental 合成一些数据,定义数据的风格样式
chain = create_data_generation_chain(model)

# 生成数据
result = chain({  # 给与一些关键词，随机生成一句话
    "fields": ['blue', 'yellow'],
    "preferences": {}
})

print(result)


# 生成一些结构化的数据，5个步骤
# 1.定义数据模型
class MedicalBilling(BaseModel):
    patient_id: int
    patient_name: str
    diagnosis_code: str
    procedure_code: str
    total_charge: float
    insurance_claim_amount: float


# 2.提供一些样例数据给AI
examples = [
    {
        "example": "Patient ID: 123456, Patient Name: zhanghui, Diagnosis Code: J20.9, Procedure Code: 123, Total Charge: 20, "
    },
    {
        "example": "Patient ID: 123456, Patient Name: zhanghui, Diagnosis Code: J20.9, Procedure Code: 123, Total Charge: 20, "
    },
]

# 3.创建一个提示模板，用来知道AI生成符合规定的数据
openai_template = PromptTemplate(input_variable=['example'], template="{example}")

prompt_template = FewShotPromptTemplate(
    prefix=SYNTHETIC_FEW_SHOT_PREFIX,
    subffix=SYNTHETIC_FEW_SHOT_SUFFIX,
    examples=examples,
    example_prompt=openai_template,
    input_variables=['subject', 'extra']
)

# 结构化数据生成器
generator = create_openai_data_generator(
    output_schema=MedicalBilling,
    llm=model,
    prompt=prompt_template
)

# 5.调用生成器
result = generator.generate(
    subject='医疗账单',  # 指定生成数据的主题
    extra='名字可以是随机的，最好使用比较生僻的人名',
    runs=10  # 指定生成数据的数量
)

print(result)
